import gradio as gr
from openai import OpenAI
import pandas as pd
import json

# Initialize the OpenAI client
client = OpenAI(
    api_key="sk-ab53207bdefa4973a739ec64ea5b9d99",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

# 数据解析函数
def parse_excel(file):
    file_path = file.name
    xls = pd.ExcelFile(file_path)
    df = None
    possible_header_rows = [0, 1, 2]
    found_columns = False
    actual_header_row = 0

    required_cols_keywords = {
        '监控时间': ['监控时间', '时间'],
        '化学需氧量': ['化学需氧量', 'COD'],
        '氨氮': ['氨氮', 'NH3-N'],
        '总氮': ['总氮', 'TN'],
        '总磷': ['总磷', 'TP']
    }

    for header_idx in possible_header_rows:
        try:
            temp_df = pd.read_excel(xls, header=header_idx)
            present_cols_count = 0
            for key_list in required_cols_keywords.values():
                if any(any(keyword.lower() in str(col).lower() for keyword in key_list) for col in temp_df.columns):
                    present_cols_count += 1
            if present_cols_count >= len(required_cols_keywords) - 1:
                df = temp_df
                actual_header_row = header_idx
                found_columns = True
                break
        except Exception:
            continue

    if not found_columns or df is None:
        try:
            df = pd.read_excel(file_path, header=[0, 1])
        except Exception:
            try:
                df = pd.read_excel(file_path, header=0)
            except Exception as e2:
                print('错误', f'无法解析Excel文件表头。请确保文件格式正确。\n错误: {e2}')
                return

    if isinstance(df.columns, pd.MultiIndex):
        df.columns = ['_'.join(map(str, col)).strip().replace('Unnamed: \d+_level_\d+_*', '', regex=True) for col in df.columns.values]
        df.columns = [col.strip('_') for col in df.columns]
    else:
        df.columns = [str(col).strip() for col in df.columns]

    df.rename(columns=lambda c: c.replace('\n', '').replace(' ', ''), inplace=True)  # Clean column names

    extracted_cols_mapping = {}
    for display_name, keywords in required_cols_keywords.items():
        found_col = None
        for col_name in df.columns:
            # More robust matching for keywords within column names
            if any(keyword.lower() in col_name.lower().replace('(毫克/升)', '').replace('mg/L', '').strip() for keyword in keywords):
                found_col = col_name
                break
        if found_col:
            extracted_cols_mapping[display_name] = found_col
        # else:
        # QMessageBox.warning(self, '列缺失警告', f'未能找到与 "{display_name}" 相关的列。该列数据将为空。')

    # Reconstruct DataFrame with only the mapped columns and standard names
    final_df_data = {}
    for display_name, actual_col_name in extracted_cols_mapping.items():
        if actual_col_name in df:
            final_df_data[display_name] = df[actual_col_name]

    final_df = pd.DataFrame(final_df_data)

    # Format time and prepare for output structure
    time_col_name = '监控时间'
    if time_col_name in final_df:
        def format_time(time_val):
            if pd.isna(time_val): return None  # Handle NaN before processing
            if isinstance(time_val, str):
                if len(time_val) == 13: return time_val + ':00'
                if len(time_val) == 10: return time_val + ' 00:00'
                # Attempt to parse other common date/time string formats if necessary
                try:
                    return pd.to_datetime(time_val).strftime('%Y-%m-%d %H:%M')
                except ValueError:
                    return str(time_val)  # Fallback for unparseable strings
            elif isinstance(time_val, pd.Timestamp):
                return time_val.strftime('%Y-%m-%d %H:%M')
            # Handle numeric/excel date serials if necessary (requires knowing epoch, e.g. for Excel)
            # elif isinstance(time_val, (int, float)):
            #   return pd.to_datetime('1899-12-30') + pd.to_timedelta(time_val, 'D') # Example for Excel serial
            return str(time_val)

        final_df[time_col_name] = final_df[time_col_name].apply(format_time)

    # Drop rows where '监控时间' is NaN after formatting, as they are key
    final_df.dropna(subset=[time_col_name], inplace=True)

    # Convert to list of dictionaries, one dict per time point (row)
    output_list = []
    for index, row in final_df.iterrows():
        row_dict = {}
        for col_name, value in row.items():
            # Ensure all values are strings for consistent JSON-like output, handle NaN
            row_dict[col_name] = str(value) if pd.notna(value) else ""
        output_list.append(row_dict)

    # Format as a JSON-like string for display and download
    converted_data_text = json.dumps(output_list, ensure_ascii=False)
    return converted_data_text

# Define the streaming response function
def chat_stream(input_text, history):
    # 不使用历史记录，每次都只发送当前输入
    promt = f"请你基于提供的数据```{input_text}```，分析每小时的污水监测曲线，找出不符合规律的异常数据。污水中包含多种物质和监测数据，不同监测数据之间没有必然联系。已知污水处理厂偶尔会通过更换采样水的方式影响监测结果，从而使监测数据出现异常变化。请你分析日常监测曲线，找出污水处理后监测数据的规律，并找出不符合规律的异常数据和时间段。请以时间维度列出日志数据、问题数据、受影响的可能性和概率。请注意，你的回答只分析结果，不包含任何代码。"

    messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},
                {'role': 'user', 'content': promt}]

    stream = client.chat.completions.create(
        model="qwen-plus",
        messages=messages,
        stream=True
    )
    partial_message = ""
    new_history = [[input_text, None]]
    for chunk in stream:
        if chunk.choices[0].delta.content is not None:
            partial_message += chunk.choices[0].delta.content
            new_history[-1][1] = partial_message
            yield new_history

# 自定义 CSS 样式
css = """
.chatbox { min-height: 400px !important; }
.clear-button { margin-top: 10px; }
"""

# Create the Gradio interface
with gr.Blocks(css=css) as demo:
    chatbot = gr.Chatbot(elem_classes="chatbox")
    file_input = gr.File(file_types=['.xlsx'], label="上传Excel文件")
    parse_and_send_btn = gr.Button("解析并发送", elem_classes="parse-and-send-button")
    clear = gr.Button("清除", elem_classes="clear-button")

    def parse_and_send(file, history):
        if file is None:
            return []  # 不保留历史记录，返回空列表
        parsed_data = parse_excel(file)
        if parsed_data is None:
            return []  # 不保留历史记录，返回空列表
        new_history = [[parsed_data, None]]
        return new_history

    parse_and_send_btn.click(parse_and_send, [file_input, chatbot], chatbot).then(
        chat_stream, [file_input, chatbot], chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7060)
